# -*- coding: utf-8 -*-

# @Time    : 2019/1/31 20:50
# @Author  : Chen
# @File    : model.py
# @Software: PyCharm

import tensorflow as tf

# 网络结构定义（一个简单的卷积神经网络，卷积+池化层x2，全连接层x2，最后一个softmax层做分类。 ）
def inference(images, batch_size, n_classes):
    # 输入参数：images, image batch, 4D tensor, tf.float32, [batch_size, width, height, hannels]
    # 返回参数：logits, float, [batch_size, n_classes]

    # 卷积层1（16个3x3的卷积核（3通道），padding=’SAME’，表示padding后卷积的图与原图尺寸一致，激活函数relu()）
    with tf.variable_scope('conv1') as scope: # tf.variable_scope(<scope_name>) 管理传给get_variable()的变量名称的作用域，通过tf.variable_scope()指定作用域进行区分，如with tf.variable_scope("conv1")这行代码指定了第一个卷积层作用域为conv1。
        weights = tf.get_variable('weights',
                                  shape = [3, 3, 3, 16],
                                  dtype = tf.float32,
                                  initializer = tf.truncated_normal_initializer(stddev = 0.1, dtype = tf.float32)) # tf.get_variable(<name>, <shape>, <initializer>) 创建或返回给定名称的变量
        biases = tf.get_variable('biases',
                                 shape = [16],
                                 dtype = tf.float32,
                                 initializer = tf.constant_initializer(0.1))
        conv = tf.nn.conv2d(images, weights,strides = [1, 1, 1, 1], padding = 'SAME')
        pre_activation = tf.nn.bias_add(conv, biases)
        conv1 = tf.nn.relu(pre_activation, name = scope.name)

    # 池化层1（3x3最大池化，步长strides为2，池化后执行lrn()操作，局部响应归一化，对训练有利。）
    with tf.variable_scope('pooling1_lrn') as scope:
        pool1 = tf.nn.max_pool(conv1, ksize = [1, 3, 3, 1], strides = [1, 2, 2, 1], padding = 'SAME', name = 'pooling1')
        norm1 = tf.nn.lrn(pool1, depth_radius = 4, bias = 1.0, alpha = 0.001 / 9.0, beta = 0.75, name = 'norm1')

    # 卷积层2（16个3x3的卷积核（16通道），padding=’SAME’，表示padding后卷积的图与原图尺寸一致，激活函数relu()）
    with tf.variable_scope('conv2') as scope:
        weights = tf.get_variable('weights',
                                  shape = [3, 3, 16, 16],
                                  dtype = tf.float32,
                                  initializer = tf.truncated_normal_initializer(stddev = 0.1, dtype = tf.float32))
        biases = tf.get_variable('biases',
                                 shape = [16],
                                 dtype = tf.float32,
                                 initializer = tf.constant_initializer(0.1))
        conv = tf.nn.conv2d(norm1, weights, strides = [1, 1, 1, 1], padding = 'SAME')
        pre_activation = tf.nn.bias_add(conv, biases)
        conv2 = tf.nn.relu(pre_activation, name = 'conv2')

    # 池化层2（3x3最大池化，步长strides为2，池化后执行lrn()操作）
    with tf.variable_scope('pooling2_lrn') as scope:
        norm2 = tf.nn.lrn(conv2, depth_radius = 4, bias = 1.0, alpha = 0.001 / 9.0, beta = 0.75, name = 'norm2')
        pool2 = tf.nn.max_pool(norm2, ksize = [1, 3, 3, 1], strides = [1, 1, 1, 1], padding = 'SAME', name = 'pooling2')

    # 全连接层3（128个神经元，将之前pool层的输出reshape成一行，激活函数relu()）
    with tf.variable_scope('local3') as scope:
        reshape = tf.reshape(pool2, shape = [batch_size, -1])
        dim = reshape.get_shape()[1].value
        weights = tf.get_variable('weights',
                                  shape = [dim, 128],
                                  dtype = tf.float32,
                                  initializer = tf.truncated_normal_initializer(stddev = 0.005, dtype = tf.float32))
        biases = tf.get_variable('biases',
                                 shape = [128],
                                 dtype = tf.float32,
                                 initializer = tf.constant_initializer(0.1))
        local3 = tf.nn.relu(tf.matmul(reshape, weights) + biases, name = scope.name)

    # 全连接层4（128个神经元，激活函数relu()）
    with tf.variable_scope('local4') as scope:
        weights = tf.get_variable('weights',
                                  shape = [128, 128],
                                  dtype = tf.float32,
                                  initializer = tf.truncated_normal_initializer(stddev = 0.005, dtype = tf.float32))
        biases = tf.get_variable('biases',
                                 shape = [128],
                                 dtype = tf.float32,
                                 initializer = tf.constant_initializer(0.1))
        local4 = tf.nn.relu(tf.matmul(local3, weights) + biases, name = 'local4')

    # Softmax回归层（将前面的FC层输出，做一个线性回归，计算出每一类的得分，在这里是2类，所以这个层输出的是两个得分。）
    with tf.variable_scope('softmax_linear') as scope:
        weights = tf.get_variable('softmax_linear',
                                  shape = [128, n_classes],
                                  dtype = tf.float32,
                                  initializer = tf.truncated_normal_initializer(stddev = 0.005, dtype = tf.float32))
        biases = tf.get_variable('biases',
                                 shape = [n_classes],
                                 dtype = tf.float32,
                                 initializer = tf.constant_initializer(0.1))
        softmax_linear = tf.add(tf.matmul(local4, weights), biases, name = 'softmax_linear')

    return softmax_linear

# loss计算（将网络计算得出的每类得分与真实值进行比较，得出一个loss损失值，这个值代表了计算值与期望值的差距。这里使用的loss函数是交叉熵。一批loss取平均数。最后调用了summary.scalar()记录下这个标量数据，在TensorBoard中进行可视化。）
def losses(logits, labels):
    # 传入参数：logits，网络计算输出值。labels，真实值，在这里是0或者1
    # 返回参数：loss，损失值

    with tf.variable_scope('loss') as scope:
        cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits = logits,
                                                                       labels = labels,
                                                                       name = 'xentropy_per_example')
        loss = tf.reduce_mean(cross_entropy, name = 'loss')
        tf.summary.scalar(scope.name + '/loss', loss)

    return loss

# loss损失值优化（目的就是让loss越小越好，使用的是AdamOptimizer优化器）
def trainning(loss, learning_rate):
    # 输入参数：loss。learning_rate，学习速率。
    # 返回参数：train_op，训练op，这个参数要输入sess.run中让模型去训练。

    with tf.name_scope('optimizer'):
        optimizer = tf.train.AdamOptimizer(learning_rate = learning_rate)
        global_step = tf.Variable(0, name = 'global_step', trainable = False)
        train_op = optimizer.minimize(loss, global_step = global_step)

    return train_op

# 评价/准确率计算（计算出平均准确率来评价这个模型，在训练过程中按批次计算（每隔N步计算一次），可以看到准确率的变换情况。）
def evaluation(logits, labels):
    # 输入参数：logits，网络计算值。labels，标签，也就是真实值，在这里是0或者1。
    # 返回参数：accuracy，当前step的平均准确率，也就是在这些batch中多少张图片被正确分类了。

    with tf.variable_scope('accuracy') as scope:
        correct = tf.nn.in_top_k(logits, labels, 1)
        correct = tf.cast(correct, tf.float16)
        accuracy = tf.reduce_mean(correct)
        tf.summary.scalar(scope.name + '/accuracy', accuracy)

    return accuracy